from torch import nn

from model.decoder.decoder import Decoder
from model.encoder.encoder import Encoder
from options import HiDDenConfiguration


class EncoderDecoder(nn.Module):

    def __init__(self, config: HiDDenConfiguration):
        super(EncoderDecoder, self).__init__()
        self.encoder = Encoder(config)
        self.decoder = Decoder(config)


    def forward(self, image, message):
        encoded_image = self.encoder(image, message)

        decoded_message = self.decoder(encoded_image)
        r